package org.sample.commons.lang;

import java.util.*;

public class TreeUtils {

    public static <T extends Node> List<T> toTree(Collection<T> nodes) {
        List<T> roots = findRoot(nodes);
        for (T root : roots)
            root.setList(findChildren(root, nodes));
        return roots;
    }

    private static <T extends Node> List<T> findRoot(Collection<T> nodes) {
        Set<Integer> ids = new HashSet<>(nodes.size());
        for (Node node : nodes)
            ids.add(node.getId());

        List<T> root = new ArrayList<>();

        Iterator<T> iterator = nodes.iterator();
        while (iterator.hasNext()) {

            T node = iterator.next();
            if (!ids.contains(node.getPid())) {

                root.add(node);
                iterator.remove();
            }
        }
        Collections.sort(root);
        return root;
    }

    private static <T extends Node> List<T> findChildren(Node root, Collection<T> nodes) {
        List<T> children = new ArrayList<>();

        Integer pid = root.getId();
        Iterator<T> iterator = nodes.iterator();
        while (iterator.hasNext()) {
            T node = iterator.next();
            if (pid == node.getPid()) {
                node.setLevel(root.getLevel() + 1);
                children.add(node);
                iterator.remove();
            }
        }

        for (Node child : children)
            child.setList(findChildren(child, nodes));

        if (children.isEmpty())
            return null;
        Collections.sort(children);
        return children;
    }

    public interface Node<T> extends Comparable<T> {
        Integer getId();

        Integer getPid();

        int getLevel();

        void setLevel(int level);

        void setList(List<T> list);
    }
}